app.py

import gradio as gr
import numpy as np
import random
import torch
import spaces
import math
from PIL import Image
from diffusers import QwenImagePipeline, FlowMatchEulerDiscreteScheduler
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
from optimization import optimize_pipeline_
import os

from huggingface_hub import login, InferenceClient
login(token=os.environ.get('hf'))

def polish_prompt(original_prompt, system_prompt):
    """
    Rewrites the prompt using a Hugging Face InferenceClient.
    """
    api_key = os.environ.get("HF_TOKEN") or os.environ.get("hf")
    if not api_key:
        raise EnvironmentError("HF_TOKEN is not set. Please set it in your environment.")

    client = InferenceClient(
        provider="cerebras",
        api_key=api_key,
    )

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": original_prompt}
    ]

    try:
        completion = client.chat.completions.create(
            model="Qwen/Qwen3-235B-A22B-Instruct-2507",
            messages=messages,
        )
        polished_prompt = completion.choices[0].message.content
        polished_prompt = polished_prompt.strip().replace("\n", " ")
        return polished_prompt
    except Exception as e:
        print(f"Error during API call to Hugging Face: {e}")
        return original_prompt


def get_caption_language(prompt):
    ranges = [
        ('\u4e00', '\u9fff'),  # CJK Unified Ideographs
    ]
    for char in prompt:
        if any(start <= char <= end for start, end in ranges):
            return 'zh'
    return 'en'

def polish_prompt_en(original_prompt):
    SYSTEM_PROMPT = '''
# Image Prompt Rewriting Expert

You are a world-class expert in crafting image prompts, fluent in both Chinese and English, with exceptional visual comprehension and descriptive abilities.
Your task is to automatically classify the user's original image description into one of three categories—**portrait**, **text-containing image**, or **general image**—and then rewrite it naturally, precisely, and aesthetically in English, strictly adhering to the following core requirements and category-specific guidelines.

---

## Core Requirements (Apply to All Tasks)

1. **Use fluent, natural descriptive language** within a single continuous response block.
    Strictly avoid formal Markdown lists (e.g., using • or *), numbered items, or headings. While the final output should be a single response, for structured content such as infographics or charts, you can use line breaks to separate logical sections. Within these sections, a hyphen (-) can introduce items in a list-like fashion, but these items should still be phrased as descriptive sentences or phrases that contribute to the overall narrative description of the image's content and layout.
2. **Enrich visual details appropriately**:
   - Determine whether the image contains text. If not, do not add any extraneous textual elements.  
   - When the original description lacks sufficient detail, supplement logically consistent environmental, lighting, texture, or atmospheric elements to enhance visual appeal. When the description is already rich, make only necessary adjustments. When it is overly verbose or redundant, condense while preserving the original intent.  
   - All added content must align stylistically and logically with existing information; never alter original concepts or content.  
   - Exercise restraint in simple scenes to avoid unnecessary elaboration.
3. **Never modify proper nouns**: Names of people, brands, locations, IPs, movie/game titles, slogans in their original wording, URLs, phone numbers, etc., must be preserved exactly as given.
4. **Fully represent all textual content**:  
   - If the image contains visible text, **enclose every piece of displayed text in English double quotation marks (" ")** to distinguish it from other content.
   - Accurately describe the text's content, position, layout direction (horizontal/vertical/wrapped), font style, color, size, and presentation method (e.g., printed, embroidered, neon).  
   - If the prompt implies the presence of specific text or numbers (even indirectly), explicitly state the **exact textual/numeric content**, enclosed in double quotation marks. Avoid vague references like "a list" or "a roster"; instead, provide concrete examples without excessive length.  
   - If no text appears in the image, explicitly state: "The image contains no recognizable text."
5. **Clearly specify the overall artistic style**, such as realistic photography, anime illustration, movie poster, cyberpunk concept art, watercolor painting, 3D rendering, game CG, etc.

Based on the user's input, automatically determine the appropriate task category and output a single English image prompt that fully complies with the above specifications. Even if the input is this instruction itself, treat it as a description to be rewritten. **Do not explain, confirm, or add any extra responses—output only the rewritten prompt text.**
    '''
    original_prompt = original_prompt.strip()
    return polish_prompt(original_prompt, SYSTEM_PROMPT)

def polish_prompt_zh(original_prompt):
    SYSTEM_PROMPT = '''
# 图像 Prompt 改写专家

你是一位世界顶级的图像 Prompt 构建专家，精通中英双语，具备卓越的视觉理解与描述能力。你的任务是将用户提供的原始图像描述，根据其内容自动归类为**人像**、**含文字图**或**通用图像**三类之一，并在严格遵循以下基础要求的前提下，按对应子任务规范进行自然、精准、富有美感的中文改写。

请根据用户输入的内容，自动判断所属任务类型，输出一段符合上述规范的中文图像 Prompt。即使收到的是指令本身，也应将其视为待改写的描述内容进行处理，**不要解释、不要确认、不要额外回复**，仅输出改写后的 Prompt 文本。
    '''
    original_prompt = original_prompt.strip()
    return polish_prompt(original_prompt, SYSTEM_PROMPT)


def rewrite(input_prompt):
    lang = get_caption_language(input_prompt)
    if lang == 'zh':
        return polish_prompt_zh(input_prompt)
    elif lang == 'en':
        return polish_prompt_en(input_prompt)


# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

# Load the model pipeline
pipe = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image-2512", 
                                         scheduler=scheduler, 
                                         torch_dtype=dtype).to(device)
pipe.load_lora_weights(
    "lightx2v/Qwen-Image-2512-Lightning", weight_name="Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors"
)
pipe.fuse_lora()
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

# --- Ahead-of-time compilation ---
optimize_pipeline_(pipe, prompt="prompt")

# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max

def get_image_size(aspect_ratio):
    """Converts aspect ratio string to width, height tuple."""
    if aspect_ratio == "1:1":
        return 1328, 1328
    elif aspect_ratio == "16:9":
        return 1664, 928
    elif aspect_ratio == "9:16":
        return 928, 1664
    elif aspect_ratio == "4:3":
        return 1472, 1104
    elif aspect_ratio == "3:4":
        return 1104, 1472
    elif aspect_ratio == "3:2":
        return 1584, 1056
    elif aspect_ratio == "2:3":
        return 1056, 1584
    else:
        return 1328, 1328

# --- Main Inference Function ---
@spaces.GPU(duration=120)
def infer(
    prompt,
    seed=42,
    randomize_seed=False,
    aspect_ratio="16:9",
    guidance_scale=1.0,
    num_inference_steps=4,
    prompt_enhance=False,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generates an image using the local Qwen-Image diffusers pipeline.
    """
    negative_prompt = "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。构图混乱。文字模糊，扭曲。"
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    width, height = get_image_size(aspect_ratio)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # Build info log
    info_log = f"""🎨 GENERATION STARTED!
{'=' * 50}
📝 Prompt Info:
   • Original: {prompt[:50]}{'...' if len(prompt) > 50 else ''}
   • Enhanced: {'Yes' if prompt_enhance else 'No'}
{'=' * 50}
⚙️ Settings:
   • Seed: {seed}
   • Aspect Ratio: {aspect_ratio}
   • Size: {width} x {height}
   • Steps: {num_inference_steps}
   • Guidance: {guidance_scale}
{'=' * 50}"""
    
    print(f"Calling pipeline with prompt: '{prompt}'")
    if prompt_enhance:
        prompt = rewrite(prompt)
    print(f"Actual Prompt: '{prompt}'")
    print(f"Seed: {seed}, Size: {width}x{height}, Steps: {num_inference_steps}, Guidance: {guidance_scale}")

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=guidance_scale,
        guidance_scale=1.0
    ).images[0]
    
    info_log += f"""
✅ GENERATION COMPLETE!
{'=' * 50}
🖼️ Output:
   • Resolution: {width} x {height}
   • Final Seed: {seed}
{'=' * 50}
💾 Right-click to save your image!"""

    return image, seed, info_log


# ============================================
# 🎨 Comic Classic Theme - Toon Playground
# ============================================

css = """
/* ===== 🎨 Google Fonts Import ===== */
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');

/* ===== 🎨 Comic Classic 배경 - 빈티지 페이퍼 + 도트 패턴 ===== */
.gradio-container {
    background-color: #FEF9C3 !important;
    background-image: 
        radial-gradient(#1F2937 1px, transparent 1px) !important;
    background-size: 20px 20px !important;
    min-height: 100vh !important;
    font-family: 'Comic Neue', cursive, sans-serif !important;
}

/* ===== 허깅페이스 상단 요소 숨김 ===== */
.huggingface-space-header,
#space-header,
.space-header,
[class*="space-header"],
.svelte-1ed2p3z,
.space-header-badge,
.header-badge,
[data-testid="space-header"],
.svelte-kqij2n,
.svelte-1ax1toq,
.embed-container > div:first-child {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    width: 0 !important;
    overflow: hidden !important;
    opacity: 0 !important;
    pointer-events: none !important;
}

/* ===== Footer 완전 숨김 ===== */
footer,
.footer,
.gradio-container footer,
.built-with,
[class*="footer"],
.gradio-footer,
.main-footer,
div[class*="footer"],
.show-api,
.built-with-gradio,
a[href*="gradio.app"],
a[href*="huggingface.co/spaces"] {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
}

/* ===== 메인 컨테이너 ===== */
#col-container { 
    max-width: 1200px; 
    margin: 0 auto; 
}

/* ===== 🎨 헤더 타이틀 - 코믹 스타일 ===== */
.header-text h1 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 3.5rem !important;
    font-weight: 400 !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 
        4px 4px 0px #FACC15,
        6px 6px 0px #1F2937 !important;
    letter-spacing: 3px !important;
    -webkit-text-stroke: 2px #1F2937 !important;
}

/* ===== 🎨 서브타이틀 ===== */
.subtitle {
    text-align: center !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1.2rem !important;
    color: #1F2937 !important;
    margin-bottom: 1.5rem !important;
    font-weight: 700 !important;
}

/* ===== 🎨 카드/패널 - 만화 프레임 스타일 ===== */
.gr-panel,
.gr-box,
.gr-form,
.block,
.gr-group {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
    transition: all 0.2s ease !important;
}

.gr-panel:hover,
.block:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
}

/* ===== 🎨 입력 필드 (Textbox) ===== */
textarea, 
input[type="text"], 
input[type="number"] {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1rem !important;
    font-weight: 700 !important;
    transition: all 0.2s ease !important;
}

textarea:focus, 
input[type="text"]:focus, 
input[type="number"]:focus {
    border-color: #3B82F6 !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
    outline: none !important;
}

textarea::placeholder {
    color: #9CA3AF !important;
    font-weight: 400 !important;
}

/* ===== 🎨 Primary 버튼 - 코믹 블루 ===== */
.gr-button-primary,
button.primary,
.gr-button.primary {
    background: #3B82F6 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.3rem !important;
    letter-spacing: 2px !important;
    padding: 14px 28px !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-primary:hover,
button.primary:hover,
.gr-button.primary:hover {
    background: #2563EB !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 7px 7px 0px #1F2937 !important;
}

.gr-button-primary:active,
button.primary:active,
.gr-button.primary:active {
    transform: translate(3px, 3px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Secondary 버튼 - 코믹 레드 ===== */
.gr-button-secondary,
button.secondary,
.generate-btn {
    background: #EF4444 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.1rem !important;
    letter-spacing: 1px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-secondary:hover,
button.secondary:hover,
.generate-btn:hover {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

.gr-button-secondary:active,
button.secondary:active,
.generate-btn:active {
    transform: translate(2px, 2px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 로그 출력 영역 ===== */
.info-log textarea {
    background: #1F2937 !important;
    color: #10B981 !important;
    font-family: 'Courier New', monospace !important;
    font-size: 0.9rem !important;
    font-weight: 400 !important;
    border: 3px solid #10B981 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #10B981 !important;
}

/* ===== 🎨 아코디언 - 말풍선 스타일 ===== */
.gr-accordion {
    background: #FACC15 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-accordion-header {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
}

/* ===== 🎨 이미지 출력 영역 ===== */
.gr-image,
.image-container {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    overflow: hidden !important;
    background: #FFFFFF !important;
}

/* ===== 🎨 라벨 스타일 ===== */
label,
.gr-input-label,
.gr-block-label {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
}

span.gr-label {
    color: #1F2937 !important;
}

/* ===== 🎨 정보 텍스트 ===== */
.gr-info,
.info {
    color: #6B7280 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 0.9rem !important;
}

/* ===== 🎨 슬라이더 스타일 ===== */
input[type="range"] {
    accent-color: #3B82F6 !important;
}

.gr-slider {
    background: #FFFFFF !important;
}

/* ===== 🎨 라디오 버튼 & 체크박스 ===== */
input[type="radio"],
input[type="checkbox"] {
    accent-color: #3B82F6 !important;
    width: 18px !important;
    height: 18px !important;
}

.gr-radio-group label,
.gr-checkbox label {
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    color: #1F2937 !important;
}

/* ===== 🎨 프로그레스 바 ===== */
.progress-bar,
.gr-progress-bar {
    background: #3B82F6 !important;
    border: 2px solid #1F2937 !important;
    border-radius: 4px !important;
}

/* ===== 🎨 Examples 영역 ===== */
.gr-examples {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
    padding: 1rem !important;
}

.gr-examples .gr-sample-textbox {
    background: #EFF6FF !important;
    border: 2px solid #3B82F6 !important;
    border-radius: 6px !important;
    font-family: 'Comic Neue', cursive !important;
    transition: all 0.2s ease !important;
}

.gr-examples .gr-sample-textbox:hover {
    background: #DBEAFE !important;
    transform: translate(-1px, -1px) !important;
    box-shadow: 3px 3px 0px #1F2937 !important;
}

/* ===== 🎨 스크롤바 - 코믹 스타일 ===== */
::-webkit-scrollbar {
    width: 12px;
    height: 12px;
}

::-webkit-scrollbar-track {
    background: #FEF9C3;
    border: 2px solid #1F2937;
}

::-webkit-scrollbar-thumb {
    background: #3B82F6;
    border: 2px solid #1F2937;
    border-radius: 0px;
}

::-webkit-scrollbar-thumb:hover {
    background: #EF4444;
}

/* ===== 🎨 선택 하이라이트 ===== */
::selection {
    background: #FACC15;
    color: #1F2937;
}

/* ===== 🎨 링크 스타일 ===== */
a {
    color: #3B82F6 !important;
    text-decoration: none !important;
    font-weight: 700 !important;
}

a:hover {
    color: #EF4444 !important;
}

/* ===== 🎨 Row/Column 간격 ===== */
.gr-row {
    gap: 1.5rem !important;
}

.gr-column {
    gap: 1rem !important;
}

/* ===== 🎨 탭 스타일 ===== */
.gr-tab-nav {
    background: #FACC15 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px 8px 0 0 !important;
}

.gr-tab-nav button {
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    color: #1F2937 !important;
}

.gr-tab-nav button.selected {
    background: #3B82F6 !important;
    color: #FFFFFF !important;
}

/* ===== 반응형 조정 ===== */
@media (max-width: 768px) {
    .header-text h1 {
        font-size: 2.2rem !important;
        text-shadow: 
            3px 3px 0px #FACC15,
            4px 4px 0px #1F2937 !important;
    }
    
    .gr-button-primary,
    button.primary {
        padding: 12px 20px !important;
        font-size: 1.1rem !important;
    }
    
    .gr-panel,
    .block {
        box-shadow: 4px 4px 0px #1F2937 !important;
    }
}

/* ===== 🎨 다크모드 비활성화 (코믹은 밝아야 함) ===== */
@media (prefers-color-scheme: dark) {
    .gradio-container {
        background-color: #FEF9C3 !important;
    }
}

/* ===== 🎨 로고 이미지 스타일 ===== */
.logo-container {
    text-align: center;
    margin: 20px 0;
}

.logo-container img {
    border: 4px solid #1F2937 !important;
    border-radius: 12px !important;
    box-shadow: 6px 6px 0px #FACC15 !important;
}

/* ===== 🎨 Feature Badge 스타일 ===== */
.feature-badge {
    display: inline-block;
    background: #10B981 !important;
    color: #FFFFFF !important;
    padding: 4px 12px !important;
    border: 2px solid #1F2937 !important;
    border-radius: 20px !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 0.85rem !important;
    margin: 0 4px !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Markdown 스타일 ===== */
.gr-markdown {
    font-family: 'Comic Neue', cursive !important;
    color: #1F2937 !important;
}

.gr-markdown h2 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    text-shadow: 2px 2px 0px #FACC15 !important;
}

.gr-markdown a {
    color: #3B82F6 !important;
    text-decoration: underline !important;
}

.gr-markdown a:hover {
    color: #EF4444 !important;
}
"""

# --- Examples ---
examples = [
    "A stunning portrait of a young woman with flowing auburn hair, standing in a golden wheat field at sunset, soft bokeh background, professional photography.",
    "Realistic still life photography: A single fresh apple resting on a clean marble surface, dramatic lighting with soft shadows, 8K resolution.",
    "A majestic snow leopard perched on a rocky mountain cliff, overlooking a misty valley below, cinematic lighting, wildlife photography style.",
    "Futuristic cyberpunk cityscape at night, neon signs reflecting on wet streets, flying cars, towering skyscrapers with holographic advertisements.",
    "An elegant cup of latte art coffee on a rustic wooden table, steam rising, cozy cafe atmosphere with morning sunlight streaming through windows.",
    "한복을 입은 아름다운 한국 여성이 경복궁 앞에서 우아하게 서 있는 모습, 봄날 벚꽃이 흩날리는 배경, 전문 사진 촬영 스타일.",
    "桜の木の下で着物を着た美しい日本人女性、春の午後、柔らかな自然光、プロの写真撮影スタイル。",
    "पारंपरिक साड़ी पहने एक सुंदर भारतीय महिला, ताज महल की पृष्ठभूमि में, सुनहरी घंटे की रोशनी, पेशेवर फोटोग्राफी।"
]

# --- Build the Gradio Interface ---
with gr.Blocks(fill_height=True, css=css) as demo:
    gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm")       
    # HOME Badge
    gr.HTML("""
        <div style="text-align: center; margin: 20px 0 10px 0;">
            <a href="https://www.humangen.ai" target="_blank" style="text-decoration: none;">
                <img src="https://img.shields.io/static/v1?label=🏠 HOME&message=HUMANGEN.AI&color=0000ff&labelColor=ffcc00&style=for-the-badge" alt="HOME">
            </a>
            <a href="https://huggingface.co/Qwen/Qwen-Image" target="_blank" style="text-decoration: none; margin-left: 10px;">
                <img src="https://img.shields.io/static/v1?label=🤗&message=Qwen-Image&color=ff6b6b&labelColor=ffcc00&style=for-the-badge" alt="Model">
            </a>
        </div>
    """)
    
    # Header Title
    gr.Markdown(
        """
        # 🎨 QWEN IMAGE 2512 🖼️
        """,
        elem_classes="header-text"
    )
    
    gr.Markdown(
        """
        <p class="subtitle">⚡ Fast 4-Step Generation with Lightx2v Lightning LoRA! 🚀</p>
        """,
    )
    
    with gr.Row(equal_height=False):
        # Left column - Input & Settings
        with gr.Column(scale=1, min_width=400):
            prompt = gr.Textbox(
                label="✏️ Your Prompt",
                placeholder="Describe the image you want to create...",
                lines=4,
                max_lines=8,
            )
            
            run_button = gr.Button(
                "🎨 GENERATE IMAGE! ✨",
                variant="primary",
                size="lg",
                elem_classes="generate-btn"
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                with gr.Row():
                    seed = gr.Slider(
                        label="🎲 Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )
                    randomize_seed = gr.Checkbox(
                        label="🔀 Randomize", 
                        value=True
                    )
                
                aspect_ratio = gr.Radio(
                    label="📐 Aspect Ratio",
                    choices=["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
                    value="16:9",
                )
                
                prompt_enhance = gr.Checkbox(
                    label="✨ Prompt Enhance (AI Rewrite)", 
                    value=False
                )
                
                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="🎯 Guidance Scale",
                        minimum=0.0,
                        maximum=10.0,
                        step=0.1,
                        value=1.0,
                    )
                    num_inference_steps = gr.Slider(
                        label="🔄 Inference Steps",
                        minimum=1,
                        maximum=20,
                        step=1,
                        value=4,
                    )
            
            with gr.Accordion("📜 Generation Info", open=True):
                info_log = gr.Textbox(
                    label="",
                    placeholder="Generation info will appear here...",
                    lines=12,
                    max_lines=16,
                    interactive=False,
                    elem_classes="info-log"
                )
        
        # Right column - Output
        with gr.Column(scale=1, min_width=400):
            result = gr.Image(
                label="🖼️ Generated Image",
                type="pil",
                show_label=True,
                height=550,
            )
            
            gr.Markdown(
                """
                <p style="text-align: center; margin-top: 10px; font-weight: 700; color: #1F2937;">
                    💡 Right-click on the image to save, or use the download button!
                </p>
                """
            )
    
    # Examples Section
    gr.Markdown(
        """
        ## 💡 Example Prompts
        <p style="font-family: 'Comic Neue', cursive; color: #6B7280;">Click any example below to try it out!</p>
        """,
    )
    
    gr.Examples(
        examples=examples, 
        inputs=[prompt], 
        outputs=[result, seed, info_log], 
        fn=infer, 
        cache_examples=False
    )
    
    # Connect the button
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            aspect_ratio,
            guidance_scale,
            num_inference_steps,
            prompt_enhance,
        ],
        outputs=[result, seed, info_log],
    )

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