app.py

import os
import subprocess
import sys
import io
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
import requests
from PIL import Image
import json
import base64
from huggingface_hub import InferenceClient

subprocess.check_call([sys.executable, "-m", "pip", "install", "spaces==0.43.0"])

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

# Pre-shifted custom sigmas for 8-step turbo inference
TURBO_SIGMAS = [1.0, 0.6509, 0.4374, 0.2932, 0.1893, 0.1108, 0.0495, 0.00031]

hf_client = InferenceClient(
    api_key=os.environ.get("HF_TOKEN"),
)
VLM_MODEL = "baidu/ERNIE-4.5-VL-424B-A47B-Base-PT"

SYSTEM_PROMPT_TEXT_ONLY = """You are an expert prompt engineer for FLUX.2 by Black Forest Labs. Rewrite user prompts to be more descriptive while strictly preserving their core subject and intent.

Guidelines:
1. Structure: Keep structured inputs structured (enhance within fields). Convert natural language to detailed paragraphs.
2. Details: Add concrete visual specifics - form, scale, textures, materials, lighting (quality, direction, color), shadows, spatial relationships, and environmental context.
3. Text in Images: Put ALL text in quotation marks, matching the prompt's language. Always provide explicit quoted text for objects that would contain text in reality (signs, labels, screens, etc.) - without it, the model generates gibberish.

Output only the revised prompt and nothing else."""

SYSTEM_PROMPT_WITH_IMAGES = """You are FLUX.2 by Black Forest Labs, an image-editing expert. You convert editing requests into one concise instruction (50-80 words, ~30 for brief requests).

Rules:
- Single instruction only, no commentary
- Use clear, analytical language (avoid "whimsical," "cascading," etc.)
- Specify what changes AND what stays the same (face, lighting, composition)
- Reference actual image elements
- Turn negatives into positives ("don't change X" → "keep X")
- Make abstractions concrete ("futuristic" → "glowing cyan neon, metallic panels")
- Keep content PG-13

Output only the final instruction in plain text and nothing else."""

def remote_text_encoder(prompts):
    from gradio_client import Client
    
    client = Client("multimodalart/mistral-text-encoder")
    result = client.predict(
        prompt=prompts,
        api_name="/encode_text"
    )
    
    prompt_embeds = torch.load(result[0])
    return prompt_embeds

# Load model
repo_id = "black-forest-labs/FLUX.2-dev"

dit = Flux2Transformer2DModel.from_pretrained(
    repo_id,
    subfolder="transformer",
    torch_dtype=torch.bfloat16
)

pipe = Flux2Pipeline.from_pretrained(
    repo_id,
    text_encoder=None,
    transformer=dit,
    torch_dtype=torch.bfloat16
)

# Load the Turbo LoRA
pipe.load_lora_weights(
    "fal/FLUX.2-dev-Turbo", 
    weight_name="flux.2-turbo-lora.safetensors"
)
pipe.fuse_lora()
pipe.unload_lora_weights()

pipe.to(device)

# Pull pre-compiled Flux2 Transformer blocks from HF hub
spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")

def image_to_data_uri(img):
    buffered = io.BytesIO()
    img.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/png;base64,{img_str}"

def upsample_prompt_logic(prompt, image_list):
    try:
        if image_list and len(image_list) > 0:
            system_content = SYSTEM_PROMPT_WITH_IMAGES
            user_content = [{"type": "text", "text": prompt}]
            
            for img in image_list:
                data_uri = image_to_data_uri(img)
                user_content.append({
                    "type": "image_url",
                    "image_url": {"url": data_uri}
                })
                
            messages = [
                {"role": "system", "content": system_content},
                {"role": "user", "content": user_content}
            ]
        else:
            system_content = SYSTEM_PROMPT_TEXT_ONLY
            messages = [
                {"role": "system", "content": system_content},
                {"role": "user", "content": prompt}
            ]

        completion = hf_client.chat.completions.create(
            model=VLM_MODEL,
            messages=messages,
            max_tokens=1024
        )
        
        return completion.choices[0].message.content
    except Exception as e:
        print(f"Upsampling failed: {e}")
        return prompt

def update_dimensions_from_image(image_list):
    """Update width/height sliders based on uploaded image aspect ratio."""
    if image_list is None or len(image_list) == 0:
        return 1024, 1024
    
    img = image_list[0][0]
    img_width, img_height = img.size
    
    aspect_ratio = img_width / img_height
    
    if aspect_ratio >= 1:
        new_width = 1024
        new_height = int(1024 / aspect_ratio)
    else:
        new_height = 1024
        new_width = int(1024 * aspect_ratio)
    
    new_width = round(new_width / 8) * 8
    new_height = round(new_height / 8) * 8
    
    new_width = max(256, min(1024, new_width))
    new_height = max(256, min(1024, new_height))
    
    return new_width, new_height

def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, use_turbo, progress=gr.Progress(track_tqdm=True)):
    num_images = 0 if image_list is None else len(image_list)
    step_duration = 1 + 0.8 * num_images
    if use_turbo:
        return max(30, 8 * step_duration + 10)
    return max(65, num_inference_steps * step_duration + 10)

@spaces.GPU(duration=get_duration)
def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, use_turbo, progress=gr.Progress(track_tqdm=True)):
    prompt_embeds = prompt_embeds.to(device)
    
    generator = torch.Generator(device=device).manual_seed(seed)
    
    pipe_kwargs = {
        "prompt_embeds": prompt_embeds,
        "image": image_list,
        "guidance_scale": guidance_scale,
        "generator": generator,
        "width": width,
        "height": height,
    }
    
    if use_turbo:
        pipe_kwargs["sigmas"] = TURBO_SIGMAS
        pipe_kwargs["num_inference_steps"] = 8
    else:
        pipe_kwargs["num_inference_steps"] = num_inference_steps
    
    if progress:
        progress(0, desc="Starting generation...")
        
    image = pipe(**pipe_kwargs).images[0]
    return image

def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=8, guidance_scale=2.5, prompt_upsampling=False, use_turbo=True, progress=gr.Progress(track_tqdm=True)):
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    image_list = None
    if input_images is not None and len(input_images) > 0:
        image_list = []
        for item in input_images:
            image_list.append(item[0])

    final_prompt = prompt
    if prompt_upsampling:
        progress(0.05, desc="Upsampling prompt...")
        final_prompt = upsample_prompt_logic(prompt, image_list)
        print(f"Original Prompt: {prompt}")
        print(f"Upsampled Prompt: {final_prompt}")

    progress(0.1, desc="Encoding prompt...")
    prompt_embeds = remote_text_encoder(final_prompt)
    
    progress(0.3, desc="Waiting for GPU...")
    image = generate_image(
        prompt_embeds, 
        image_list, 
        width, 
        height, 
        num_inference_steps, 
        guidance_scale, 
        seed,
        use_turbo,
        progress
    )
    
    # 정보 로그 생성
    info_log = f"""✅ GENERATION COMPLETE!
{'=' * 50}
📝 Prompt Info:
   • Original: {prompt[:50]}{'...' if len(prompt) > 50 else ''}
   • Upsampled: {'Yes' if prompt_upsampling else 'No'}
{'=' * 50}
⚙️ Generation Settings:
   • Seed: {seed}
   • Size: {width} x {height}
   • Steps: {'8 (Turbo)' if use_turbo else num_inference_steps}
   • CFG Scale: {guidance_scale}
   • Input Images: {len(image_list) if image_list else 0}
{'=' * 50}
🚀 Mode: {'⚡ TURBO (8 steps)' if use_turbo else '🎨 Standard'}
{'=' * 50}
💾 Image ready to download!"""
    
    return image, seed, info_log

examples = [
    ["Create a vase on a table in living room, the color of the vase is a gradient of color, starting with #02eb3c color and finishing with #edfa3c. The flowers inside the vase have the color #ff0088"],
    ["Photorealistic infographic showing the complete Berlin TV Tower (Fernsehturm) from ground base to antenna tip, full vertical view with entire structure visible including concrete shaft, metallic sphere, and antenna spire. Slight upward perspective angle looking up toward the iconic sphere, perfectly centered on clean white background. Left side labels with thin horizontal connector lines: the text '368m' in extra large bold dark grey numerals (#2D3748) positioned at exactly the antenna tip with 'TOTAL HEIGHT' in small caps below. The text '207m' in extra large bold with 'TELECAFÉ' in small caps below, with connector line touching the sphere precisely at the window level. Right side label with horizontal connector line touching the sphere's equator: the text '32m' in extra large bold dark grey numerals with 'SPHERE DIAMETER' in small caps below. Bottom section arranged in three balanced columns: Left - Large text '986' in extra bold dark grey with 'STEPS' in caps below. Center - 'BERLIN TV TOWER' in bold caps with 'FERNSEHTURM' in lighter weight below. Right - 'INAUGURATED' in bold caps with 'OCTOBER 3, 1969' below. All typography in modern sans-serif font (such as Inter or Helvetica), color #2D3748, clean minimal technical diagram style. Horizontal connector lines are thin, precise, and clearly visible, touching the tower structure at exact corresponding measurement points. Professional architectural elevation drawing aesthetic with dynamic low angle perspective creating sense of height and grandeur, poster-ready infographic design with perfect visual hierarchy."],
    ["Soaking wet capybara taking shelter under a banana leaf in the rainy jungle, close up photo"],
    ["A kawaii die-cut sticker of a chubby orange cat, featuring big sparkly eyes and a happy smile with paws raised in greeting and a heart-shaped pink nose. The design should have smooth rounded lines with black outlines and soft gradient shading with pink cheeks."],
]

examples_images = [
    ["The person from image 1 is petting the cat from image 2, the bird from image 3 is next to them", ["woman1.webp", "cat_window.webp", "bird.webp"]]
]


# ============================================
# 🎨 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;
}

.subtitle-small {
    text-align: center !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1rem !important;
    color: #6B7280 !important;
    margin-bottom: 1rem !important;
    font-weight: 400 !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;
}

/* ===== 🎨 프롬프트 입력창 특별 스타일 ===== */
.prompt-input textarea {
    background: #FFFBEB !important;
    border: 4px solid #F59E0B !important;
    border-radius: 12px !important;
    font-size: 1.1rem !important;
    padding: 12px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.prompt-input textarea:focus {
    border-color: #3B82F6 !important;
    box-shadow: 6px 6px 0px #3B82F6 !important;
}

/* ===== 🎨 Primary 버튼 - 코믹 블루 ===== */
.gr-button-primary,
button.primary,
.gr-button.primary,
.generate-btn {
    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,
.generate-btn: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,
.generate-btn:active {
    transform: translate(3px, 3px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Secondary 버튼 - 코믹 레드 ===== */
.gr-button-secondary,
button.secondary {
    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 {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 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;
}

/* ===== 🎨 이미지 업로드/갤러리 영역 ===== */
.image-upload,
.gr-gallery {
    border: 4px dashed #3B82F6 !important;
    border-radius: 12px !important;
    background: #EFF6FF !important;
    transition: all 0.2s ease !important;
}

.image-upload:hover,
.gr-gallery:hover {
    border-color: #EF4444 !important;
    background: #FEF2F2 !important;
}

.gr-gallery .thumbnail-item {
    border: 3px solid #1F2937 !important;
    border-radius: 6px !important;
    transition: all 0.2s ease !important;
}

.gr-gallery .thumbnail-item:hover {
    transform: scale(1.05) !important;
    box-shadow: 4px 4px 0px #3B82F6 !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;
}

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

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

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

/* ===== 🎨 체크박스 스타일 ===== */
input[type="checkbox"] {
    accent-color: #3B82F6 !important;
    width: 20px !important;
    height: 20px !important;
    border: 2px solid #1F2937 !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;
}

/* ===== 🎨 프로그레스 바 ===== */
.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: 6px 6px 0px #1F2937 !important;
    padding: 1rem !important;
}

.gr-examples .gr-sample {
    border: 2px solid #1F2937 !important;
    border-radius: 6px !important;
    transition: all 0.2s ease !important;
}

.gr-examples .gr-sample:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
}

/* ===== 🎨 Turbo 뱃지 스타일 ===== */
.turbo-badge {
    display: inline-block;
    background: linear-gradient(135deg, #F59E0B 0%, #EF4444 100%) !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-size: 1rem !important;
    padding: 4px 12px !important;
    border: 2px solid #1F2937 !important;
    border-radius: 20px !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
    margin-left: 8px !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;
}

/* ===== 반응형 조정 ===== */
@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;
    }
}
"""

# 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>
        </div>
    """)
    
    # Header Title
    gr.Markdown(
        """
        # ⚡ FLUX.2 TURBO IMAGE GENERATOR 🎨
        """,
        elem_classes="header-text"
    )
    
    gr.Markdown(
        """
        <p class="subtitle">🚀 32B Rectified Flow Model • Generate, Edit & Combine Images in 8 Steps! ✨</p>
        <p class="subtitle-small">Powered by <a href="https://huggingface.co/black-forest-labs/FLUX.2-dev" target="_blank">FLUX.2 [dev]</a> with <a href="https://huggingface.co/fal/FLUX.2-Turbo" target="_blank">Turbo LoRA by fal</a></p>
        """,
    )
    
    with gr.Row(equal_height=False):
        # Left column - Input
        with gr.Column(scale=1, min_width=400):
            prompt = gr.Textbox(
                label="✏️ Enter Your Prompt",
                placeholder="Describe the image you want to create...",
                lines=3,
                max_lines=5,
                elem_classes="prompt-input"
            )
            
            run_button = gr.Button(
                "⚡ GENERATE IMAGE! 🎨",
                variant="primary",
                size="lg",
                elem_classes="generate-btn"
            )
            
            with gr.Accordion("🖼️ Input Images (Optional)", open=True):
                input_images = gr.Gallery(
                    label="Upload reference images for editing/combining",
                    type="pil",
                    columns=3,
                    rows=1,
                    elem_classes="image-upload"
                )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                use_turbo = gr.Checkbox(
                    label="⚡ Use Turbo Mode (8 steps)",
                    value=True,
                    info="Enable Turbo LoRA for fast 8-step generation",
                    visible=False
                )
                
                prompt_upsampling = gr.Checkbox(
                    label="🔮 Prompt Upsampling",
                    value=False,
                    info="Automatically enhance the prompt using a VLM"
                )
    
                seed = gr.Slider(
                    label="🎲 Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                
                randomize_seed = gr.Checkbox(label="🔀 Randomize seed", value=True)
                
                with gr.Row():
                    width = gr.Slider(
                        label="📐 Width",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=8,
                        value=1024,
                    )
                    
                    height = gr.Slider(
                        label="📏 Height",
                        minimum=256,
                        maximum=MAX_IMAGE_SIZE,
                        step=8,
                        value=1024,
                    )
                
                with gr.Row():
                    num_inference_steps = gr.Slider(
                        label="🔄 Inference Steps (ignored in Turbo)",
                        minimum=1,
                        maximum=100,
                        step=1,
                        value=30,
                    )
                    
                    guidance_scale = gr.Slider(
                        label="🎯 Guidance Scale",
                        minimum=0.0,
                        maximum=10.0,
                        step=0.1,
                        value=2.5,
                    )
            
            with gr.Accordion("📜 Generation Log", open=True):
                info_log = gr.Textbox(
                    label="",
                    placeholder="Enter a prompt and click generate to see info...",
                    lines=14,
                    max_lines=20,
                    interactive=False,
                    elem_classes="info-log"
                )
        
        # Right column - Output
        with gr.Column(scale=1, min_width=400):
            result = gr.Image(
                label="🖼️ Generated Image",
                show_label=True,
                height=550,
                elem_classes="result-image"
            )
            
            gr.Markdown(
                """
                <p style="text-align: center; margin-top: 15px; font-weight: 700; color: #1F2937;">
                    💡 Right-click on the image to save, or use the download button!
                </p>
                """
            )

    # Examples Section
    gr.Markdown(
        """
        <p style="text-align: center; margin: 25px 0 15px 0; font-family: 'Bangers', cursive; font-size: 1.5rem; color: #1F2937;">
            🌟 TRY THESE EXAMPLES! 🌟
        </p>
        """
    )
    
    gr.Examples(
        examples=examples,
        fn=infer,
        inputs=[prompt],
        outputs=[result, seed, info_log],
        cache_examples=True,
        cache_mode="lazy"
    )

    gr.Examples(
        examples=examples_images,
        fn=infer,
        inputs=[prompt, input_images],
        outputs=[result, seed, info_log],
        cache_examples=True,
        cache_mode="lazy"
    )

    # Auto-update dimensions when images are uploaded
    input_images.upload(
        fn=update_dimensions_from_image,
        inputs=[input_images],
        outputs=[width, height]
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling, use_turbo],
        outputs=[result, seed, info_log]
    )

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